"""
PyTorch模型检查点管理器(Checkpointer)

核心功能：
1. 自动化模型参数的保存与加载
2. 支持版本控制（计数+时间戳）
3. 提供队列机制管理历史检查点
4. 兼容单机/多GPU训练模式

典型应用场景：
- 深度学习模型训练过程中的断点续训
- 模型性能比较时的快照保存
- 生产环境模型部署的版本管理
"""

import os
import time
from os import path
from queue import Queue
import torch
from tools.paths import process


class Checkpointer(object):
    """PyTorch对象序列化管理器

    参数说明：
        name (str): 检查点名称（用于生成文件名）
        directory (str): 存储目录（默认当前目录）
        overwrite (bool): 是否覆盖旧检查点（默认False）
        verbose (bool): 是否打印操作日志（默认True）
        timestamp (bool): 文件名添加时间戳（默认False）
        add_count (bool): 文件名添加计数（默认True）
        max_queue (int): 最大保存队列长度（None表示不限制）

    文件命名规则：
        [name]_[tag]_[counter]_[timestamp].pth
    """

    def __init__(self, name, directory='.', overwrite=False, verbose=True,
                 timestamp=False, add_count=True, max_queue=None):
        """初始化检查点管理器"""
        self.name = name
        self.directory = process(directory, create=True)  # 确保目录存在
        self.directory = path.join(self.directory, 'checkpoints')  # 创建checkpoints子目录
        self.directory = process(self.directory, create=True)
        self.overwrite = overwrite
        self.timestamp = timestamp
        self.add_count = add_count
        self.verbose = verbose
        self.chkp = path.join(self.directory, ".{0}.chkp".format(self.name))  # 元数据文件

        # 加载已有检查点或初始化
        self.counter, self.filename = torch.load(self.chkp) if path.exists(self.chkp) else (0, '')
        self.save_queue = False

        # 初始化检查点队列（用于限制保存数量）
        if overwrite == False and isinstance(max_queue, int) and max_queue > 1:
            self.save_queue = True
            self._queue = Queue(max_queue)
        elif max_queue is not None:
            print('WARNING: illegal max_queue Value!.')

        self.show_save_pth_name = ''  # 最近保存路径显示

    def _say(self, line):
        """日志输出方法"""
        if self.verbose:
            print(line)

    def _make_name(self, tag=None):
        """生成检查点文件名

        参数：
            tag (str): 自定义标签

        返回：
            str: 完整文件路径
        """
        strend = ''
        if tag is not None:
            strend += '_' + str(tag)
        if self.add_count:
            strend += "_{:07d}".format(self.counter)  # 7位零填充计数
        if self.timestamp:
            strend += time.strftime("_%Y-%m-%d-%H-%M-%S")  # 时间戳
        filename = "{0}{1}.pth".format(self.name, strend)
        self.filename = path.join(self.directory, filename)
        return self.filename

    def _set_state(self, obj, state):
        """加载状态字典到对象

        支持两种模式：
        1. 标准PyTorch模块（使用load_state_dict）
        2. 普通Python对象（直接返回状态）
        """
        if hasattr(obj, 'load_state_dict'):
            obj.load_state_dict(state)
            return obj
        else:
            return state

    def _get_state(self, obj):
        """提取对象状态字典

        特殊处理：
        - 多GPU训练包装（DataParallel）
        - 普通PyTorch模块（state_dict）
        - 其他Python对象（直接返回）
        """
        if isinstance(obj, torch.nn.DataParallel):
            obj = obj.module  # 解包DataParallel
        if hasattr(obj, 'state_dict'):
            return obj.state_dict()
        else:
            return obj

    def __call__(self, obj, tag=None, *args, **kwargs):
        """快捷调用方式（等效于save）"""
        self.save(obj, tag=tag, *args, **kwargs)

    def save(self, obj, tag=None, *args, **kwargs):
        """保存对象检查点

        处理流程：
        1. 生成新文件名
        2. 管理检查点队列（删除最旧文件）
        3. 序列化对象状态
        4. 更新元数据
        """
        self.counter += 1
        old_filename = self.filename
        new_filename = self._make_name(tag)

        # 队列管理（限制历史版本数量）
        if self.save_queue is True:
            if self._queue.full() is True:
                delete_name = self._queue.get()  # 获取最旧文件名
                try:
                    os.remove(delete_name)  # 删除旧文件
                except:
                    pass
            self._queue.put(new_filename)  # 加入新文件

        # 文件覆盖处理
        if new_filename == old_filename and not self.overwrite:
            print('WARNING: Overwriting file in non overwrite mode.')
        elif self.overwrite:
            try:
                os.remove(old_filename)
            except:
                pass

        # 执行序列化
        torch.save(self._get_state(obj), new_filename, *args, **kwargs)
        torch.save((self.counter, new_filename), self.chkp)  # 保存元数据

        self.show_save_pth_name = new_filename  # 更新显示路径

    def load(self, obj=None, preprocess=None, multi_gpu=False, *args, **kwargs):
        """加载检查点

        特殊功能：
        - 支持直接加载.pth文件路径
        - 可选的预处理函数
        - 多GPU模式自动转换（添加module.前缀）
        """
        if isinstance(obj, str) and obj.split('.')[-1] == 'pth':  # 直接加载文件
            self._say("Loaded checkpoint: {0}".format(obj))
            obj = torch.load(obj)
            # 修改checkpointer.py第171行为：   torch版本不一样这里的weights_only就不一样  add by likang
            # obj = torch.load(obj, weights_only=True)  # 显式启用安全加载
        elif self.counter > 0 and obj is None:  # 加载最新检查点
            loaded = torch.load(self.filename, *args, **kwargs)
            if preprocess is not None:  # 执行预处理
                loaded = preprocess(loaded)
            obj = self._set_state(obj, loaded)
            self._say("Loaded {0} checkpoint: {1}".format(self.name, self.filename))

        # 多GPU模式转换
        if multi_gpu is True:
            from collections import OrderedDict
            multi_gpu_obj = OrderedDict()
            for k, v in obj.items():
                name = 'module.' + k  # 添加DataParallel前缀
                multi_gpu_obj[name] = v
            obj = multi_gpu_obj

        return obj